package penguin;

import penguin.masking.Interceptors;
import penguin.masking.Logger;
import penguin.masking.MaskingContext;

import java.lang.instrument.Instrumentation;
import java.lang.instrument.UnmodifiableClassException;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

/**
 * 代理类
 * <p>负责完成代理到目标jvm</p>
 *
 * @author wuxh
 * @version 1.0.0
 */
public class DataMaskingAgent {

    public static final String COMMAND_DETACH = "uninstall";
    public static final String COMMAND_STOP = "stop";
    public static final String COMMAND_INSTALL = "install";

    //-javaagent
    public static void premain(String agentArgs, Instrumentation inst) throws UnmodifiableClassException {
        agent(agentArgs, inst);
    }

    //attach api
    public static void agentmain(String agentArgs, Instrumentation inst) throws UnmodifiableClassException {
        agent(agentArgs, inst);
    }

    private static void agent(String agentArgs, Instrumentation inst) throws UnmodifiableClassException {
        String command = agentArgs;
        final MaskingContext context = MaskingContext.context();
        System.out.println("COMMAND:" + agentArgs);
        if (COMMAND_DETACH.equals(command)) {//卸载
            if (context.isUninstall()) {
                throw new RuntimeException("程序未安装，不能执行卸载命令！");
            }
            Interceptors interceptors = context.getInterceptors();
            inst.removeTransformer(interceptors);
            context.uninstall();
            return;
        } else if (COMMAND_STOP.equals(command)) {//停止
            if (!context.isRunning()) {
                throw new RuntimeException("程序未运行，不能执行停止命令！");
            }
            context.stop();
            return;
        }
        if (!COMMAND_INSTALL.equals(command)) {
            return;
        }
        if (!context.isUninstall()) {
            throw new RuntimeException("程序在运行中，不能执行安装命令！");
        }
        context.start();
        Set<String> notLoadClass = new HashSet<>(context.getInterceptors().reloadClz());
        Class classes[] = inst.getAllLoadedClasses();
        List<Class> retransformClasses = new ArrayList<>();
        for (int i = 0; i < classes.length; i++) {
            if (context.getInterceptors().containsClass(classes[i].getName())) {//当前包的类
                retransformClasses.add(classes[i]);
                notLoadClass.remove(classes[i].getName());
                Logger.debug("Reloading: " + classes[i].getName());
            }
        }
        //触发类的加载
        //以-javaagent执行时，类未加载进去。如果是有依赖的类，不触发ClassFileTransformer，导致字节码增强不了
        //详情看Instrumentation.addTransformer
        for (String loadClass : notLoadClass) {
            try {
                Class<?> aClass = Class.forName(loadClass);
                retransformClasses.add(aClass);
            } catch (ClassNotFoundException e) {
                //do nothing
                e.printStackTrace();
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
        inst.addTransformer(context.getInterceptors(), true);
        if (!retransformClasses.isEmpty()) {
            Class[] _classes = retransformClasses.toArray(new Class[retransformClasses.size()]);
            inst.retransformClasses(_classes);
        }
    }
}
